import sqlite3
from datetime import datetime

# 数据库连接函数
def get_db_connection():
    conn = sqlite3.connect('admin.db')
    conn.row_factory = sqlite3.Row
    return conn

class User:
    def __init__(self, username, password, email=None, nickname=None, status=1):
        self.id = None
        self.username = username
        self.password = password
        self.email = email
        self.nickname = nickname
        self.status = status
        self.created_at = datetime.now()
        self.updated_at = datetime.now()

class Role:
    def __init__(self, name, description=None, status=1):
        self.id = None
        self.name = name
        self.description = description
        self.status = status
        self.created_at = datetime.now()
        self.updated_at = datetime.now()

class Permission:
    def __init__(self, name, code, description=None, category=None):
        self.id = None
        self.name = name
        self.code = code
        self.description = description
        self.category = category
        self.created_at = datetime.now()

class Menu:
    def __init__(self, name, parent_id=0, icon=None, path=None, component=None, 
                 sort_order=0, status=1, permission_id=None):
        self.id = None
        self.parent_id = parent_id
        self.name = name
        self.icon = icon
        self.path = path
        self.component = component
        self.sort_order = sort_order
        self.status = status
        self.permission_id = permission_id
        self.created_at = datetime.now()
        self.updated_at = datetime.now()

class OperationLog:
    def __init__(self, user_id, username, operation, method=None, url=None, ip=None, 
                 user_agent=None, params=None, result=None, duration=None):
        self.id = None
        self.user_id = user_id
        self.username = username
        self.operation = operation
        self.method = method
        self.url = url
        self.ip = ip
        self.user_agent = user_agent
        self.params = params
        self.result = result
        self.duration = duration
        self.created_at = datetime.now()

class LoginLog:
    def __init__(self, user_id, username, ip=None, user_agent=None, status=None, message=None):
        self.id = None
        self.user_id = user_id
        self.username = username
        self.ip = ip
        self.user_agent = user_agent
        self.status = status
        self.message = message
        self.created_at = datetime.now()

class SystemConfig:
    def __init__(self, key, value, description=None, category=None):
        self.id = None
        self.key = key
        self.value = value
        self.description = description
        self.category = category
        self.created_at = datetime.now()
        self.updated_at = datetime.now()

class DataDictionary:
    def __init__(self, code, name, category=None, description=None, sort_order=0, status=1):
        self.id = None
        self.code = code
        self.name = name
        self.category = category
        self.description = description
        self.sort_order = sort_order
        self.status = status
        self.created_at = datetime.now()
        self.updated_at = datetime.now()

class DictionaryItem:
    def __init__(self, dictionary_id, code, label, value=None, description=None, sort_order=0, status=1):
        self.id = None
        self.dictionary_id = dictionary_id
        self.code = code
        self.label = label
        self.value = value
        self.description = description
        self.sort_order = sort_order
        self.status = status
        self.created_at = datetime.now()
        self.updated_at = datetime.now()

def init_db():
    conn = get_db_connection()
    
    # 创建用户表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username VARCHAR(50) UNIQUE NOT NULL,
            password VARCHAR(255) NOT NULL,
            email VARCHAR(100) UNIQUE,
            nickname VARCHAR(100),
            status TINYINT DEFAULT 1,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建角色表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS roles (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name VARCHAR(50) UNIQUE NOT NULL,
            description VARCHAR(255),
            status TINYINT DEFAULT 1,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建权限表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS permissions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name VARCHAR(100) NOT NULL,
            code VARCHAR(100) UNIQUE NOT NULL,
            description VARCHAR(255),
            category VARCHAR(50),
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建菜单表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS menus (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            parent_id INTEGER DEFAULT 0,
            name VARCHAR(100) NOT NULL,
            icon VARCHAR(100),
            path VARCHAR(255),
            component VARCHAR(255),
            sort_order INT DEFAULT 0,
            status TINYINT DEFAULT 1,
            permission_id INTEGER,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建用户角色关联表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS user_roles (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            role_id INTEGER NOT NULL,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建角色权限关联表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS role_permissions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            role_id INTEGER NOT NULL,
            permission_id INTEGER NOT NULL,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建操作日志表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS operation_logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER,
            username VARCHAR(50),
            operation VARCHAR(100) NOT NULL,
            method VARCHAR(10),
            url VARCHAR(255),
            ip VARCHAR(45),
            user_agent TEXT,
            params TEXT,
            result TEXT,
            duration INT,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建登录日志表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS login_logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER,
            username VARCHAR(50),
            ip VARCHAR(45),
            user_agent TEXT,
            status TINYINT,
            message VARCHAR(255),
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建系统配置表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS system_configs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            key VARCHAR(100) UNIQUE NOT NULL,
            value TEXT,
            description VARCHAR(255),
            category VARCHAR(50),
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建数据字典表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS data_dictionaries (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            code VARCHAR(50) NOT NULL,
            name VARCHAR(100) NOT NULL,
            category VARCHAR(50),
            description VARCHAR(255),
            sort_order INT DEFAULT 0,
            status TINYINT DEFAULT 1,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    # 创建字典项表
    conn.execute('''
        CREATE TABLE IF NOT EXISTS dictionary_items (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            dictionary_id INTEGER NOT NULL,
            code VARCHAR(50) NOT NULL,
            label VARCHAR(100) NOT NULL,
            value VARCHAR(255),
            description VARCHAR(255),
            sort_order INT DEFAULT 0,
            status TINYINT DEFAULT 1,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    
    conn.commit()
    conn.close()

# 用户相关操作
def create_user(username, password, email=None, nickname=None, status=1):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO users (username, password, email, nickname, status) VALUES (?, ?, ?, ?, ?)',
        (username, password, email, nickname, status)
    )
    user_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return user_id

def get_all_users():
    conn = get_db_connection()
    users = conn.execute('SELECT * FROM users').fetchall()
    conn.close()
    return users

def get_user_by_id(user_id):
    conn = get_db_connection()
    user = conn.execute('SELECT * FROM users WHERE id = ?', (user_id,)).fetchone()
    conn.close()
    return user

# 一个有注入漏洞的函数（正常请使用上方的get_user_by_id方法）
def get_user_by_id_with_sql_inject(user_id):
    conn = get_db_connection()
    user = conn.execute('SELECT * FROM users WHERE id =  ' + user_id).fetchone()
    conn.close()
    return user

def get_user_by_username(username):
    conn = get_db_connection()
    user = conn.execute('SELECT * FROM users WHERE username = ?', (username,)).fetchone()
    conn.close()
    return user

# 用户角色相关操作
def get_user_roles(user_id):
    conn = get_db_connection()
    roles = conn.execute('''
        SELECT r.* FROM roles r
        JOIN user_roles ur ON r.id = ur.role_id
        WHERE ur.user_id = ?
    ''', (user_id,)).fetchall()
    conn.close()
    return roles

def assign_role_to_user(user_id, role_id):
    conn = get_db_connection()
    cursor = conn.cursor()
    try:
        cursor.execute(
            'INSERT INTO user_roles (user_id, role_id) VALUES (?, ?)',
            (user_id, role_id)
        )
        conn.commit()
        return True
    except sqlite3.IntegrityError:
        return False
    finally:
        conn.close()

def remove_role_from_user(user_id, role_id):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'DELETE FROM user_roles WHERE user_id = ? AND role_id = ?',
        (user_id, role_id)
    )
    conn.commit()
    conn.close()

# 角色相关操作
def create_role(name, description=None, status=1):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO roles (name, description, status) VALUES (?, ?, ?)',
        (name, description, status)
    )
    role_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return role_id

def get_all_roles():
    conn = get_db_connection()
    roles = conn.execute('SELECT * FROM roles').fetchall()
    conn.close()
    return roles

def get_role_by_id(role_id):
    conn = get_db_connection()
    role = conn.execute('SELECT * FROM roles WHERE id = ?', (role_id,)).fetchone()
    conn.close()
    return role

# 角色权限相关操作
def get_role_permissions(role_id):
    conn = get_db_connection()
    permissions = conn.execute('''
        SELECT p.* FROM permissions p
        JOIN role_permissions rp ON p.id = rp.permission_id
        WHERE rp.role_id = ?
    ''', (role_id,)).fetchall()
    conn.close()
    return permissions

def assign_permission_to_role(role_id, permission_id):
    conn = get_db_connection()
    cursor = conn.cursor()
    try:
        cursor.execute(
            'INSERT INTO role_permissions (role_id, permission_id) VALUES (?, ?)',
            (role_id, permission_id)
        )
        conn.commit()
        return True
    except sqlite3.IntegrityError:
        return False
    finally:
        conn.close()

def remove_permission_from_role(role_id, permission_id):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'DELETE FROM role_permissions WHERE role_id = ? AND permission_id = ?',
        (role_id, permission_id)
    )
    conn.commit()
    conn.close()

# 权限相关操作
def create_permission(name, code, description=None, category=None):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO permissions (name, code, description, category) VALUES (?, ?, ?, ?)',
        (name, code, description, category)
    )
    permission_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return permission_id

def get_all_permissions():
    conn = get_db_connection()
    permissions = conn.execute('SELECT * FROM permissions').fetchall()
    conn.close()
    return permissions

def get_permission_by_code(code):
    conn = get_db_connection()
    permission = conn.execute('SELECT * FROM permissions WHERE code = ?', (code,)).fetchone()
    conn.close()
    return permission

# 菜单相关操作
def create_menu(name, parent_id=0, icon=None, path=None, component=None, 
                sort_order=0, status=1, permission_id=None):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO menus (name, parent_id, icon, path, component, sort_order, status, permission_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?)',
        (name, parent_id, icon, path, component, sort_order, status, permission_id)
    )
    menu_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return menu_id

def get_all_menus():
    conn = get_db_connection()
    menus = conn.execute('SELECT * FROM menus').fetchall()
    conn.close()
    return menus

def get_menu_by_id(menu_id):
    conn = get_db_connection()
    menu = conn.execute('SELECT * FROM menus WHERE id = ?', (menu_id,)).fetchone()
    conn.close()
    return menu

def get_menus_by_user_id(user_id):
    conn = get_db_connection()
    # 获取用户所有角色
    user_roles = conn.execute('''
        SELECT r.id FROM roles r
        JOIN user_roles ur ON r.id = ur.role_id
        WHERE ur.user_id = ? AND r.status = 1
    ''', (user_id,)).fetchall()
    
    if not user_roles:
        # 如果用户没有角色，则只显示无权限要求的菜单
        menus = conn.execute('SELECT * FROM menus WHERE permission_id IS NULL AND status = 1 ORDER BY sort_order').fetchall()
    else:
        # 获取用户角色对应的权限
        role_ids = [role[0] for role in user_roles]
        placeholders = ','.join('?' * len(role_ids))
        permissions = conn.execute(f'''
            SELECT DISTINCT p.id FROM permissions p
            JOIN role_permissions rp ON p.id = rp.permission_id
            WHERE rp.role_id IN ({placeholders})
        ''', role_ids).fetchall()
        
        permission_ids = [perm[0] for perm in permissions]
        
        if permission_ids:
            # 显示有权限的菜单和无权限要求的菜单
            perm_placeholders = ','.join('?' * len(permission_ids))
            menus = conn.execute(f'''
                SELECT * FROM menus 
                WHERE (permission_id IN ({perm_placeholders}) OR permission_id IS NULL) 
                AND status = 1 
                ORDER BY sort_order
            ''', permission_ids).fetchall()
        else:
            # 如果角色没有权限，则只显示无权限要求的菜单
            menus = conn.execute('SELECT * FROM menus WHERE permission_id IS NULL AND status = 1 ORDER BY sort_order').fetchall()
    
    conn.close()
    return menus

# 操作日志相关操作
def create_operation_log(user_id, username, operation, method=None, url=None, ip=None, 
                        user_agent=None, params=None, result=None, duration=None):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        '''INSERT INTO operation_logs 
           (user_id, username, operation, method, url, ip, user_agent, params, result, duration) 
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
        (user_id, username, operation, method, url, ip, user_agent, params, result, duration)
    )
    log_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return log_id

def get_all_operation_logs():
    conn = get_db_connection()
    logs = conn.execute('SELECT * FROM operation_logs ORDER BY created_at DESC').fetchall()
    conn.close()
    return logs

# 登录日志相关操作
def create_login_log(user_id, username, ip=None, user_agent=None, status=None, message=None):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO login_logs (user_id, username, ip, user_agent, status, message) VALUES (?, ?, ?, ?, ?, ?)',
        (user_id, username, ip, user_agent, status, message)
    )
    log_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return log_id

def get_all_login_logs():
    conn = get_db_connection()
    logs = conn.execute('SELECT * FROM login_logs ORDER BY created_at DESC').fetchall()
    conn.close()
    return logs

# 系统配置相关操作
def create_system_config(key, value, description=None, category=None):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO system_configs (key, value, description, category) VALUES (?, ?, ?, ?)',
        (key, value, description, category)
    )
    config_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return config_id

def get_all_system_configs():
    conn = get_db_connection()
    configs = conn.execute('SELECT * FROM system_configs').fetchall()
    conn.close()
    return configs

def get_system_config_by_key(key):
    conn = get_db_connection()
    config = conn.execute('SELECT * FROM system_configs WHERE key = ?', (key,)).fetchone()
    conn.close()
    return config

def update_system_config(key, value, description=None, category=None):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'UPDATE system_configs SET value = ?, description = ?, category = ? WHERE key = ?',
        (value, description, category, key)
    )
    conn.commit()
    conn.close()

# 数据字典相关操作
def create_data_dictionary(code, name, category=None, description=None, sort_order=0, status=1):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO data_dictionaries (code, name, category, description, sort_order, status) VALUES (?, ?, ?, ?, ?, ?)',
        (code, name, category, description, sort_order, status)
    )
    dictionary_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return dictionary_id

def get_all_data_dictionaries():
    conn = get_db_connection()
    dictionaries = conn.execute('SELECT * FROM data_dictionaries').fetchall()
    conn.close()
    return dictionaries

def get_data_dictionary_by_code(code):
    conn = get_db_connection()
    dictionary = conn.execute('SELECT * FROM data_dictionaries WHERE code = ?', (code,)).fetchone()
    conn.close()
    return dictionary

# 字典项相关操作
def create_dictionary_item(dictionary_id, code, label, value=None, description=None, sort_order=0, status=1):
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        'INSERT INTO dictionary_items (dictionary_id, code, label, value, description, sort_order, status) VALUES (?, ?, ?, ?, ?, ?, ?)',
        (dictionary_id, code, label, value, description, sort_order, status)
    )
    item_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return item_id

def get_dictionary_items_by_dictionary_id(dictionary_id):
    conn = get_db_connection()
    items = conn.execute('SELECT * FROM dictionary_items WHERE dictionary_id = ?', (dictionary_id,)).fetchall()
    conn.close()
    return items

if __name__ == '__main__':
    init_db()
    print("Database initialized successfully.")